import getopt
import sys
import oss2
import pymysql
import redis
import time
import json
import datetime
from PIL import Image
import io, os
import requests
import traceback
import base64
import random
from urllib.parse import urlparse

env = "prod"

opts, args = getopt.getopt(sys.argv[1:], "", "env=")
for opt, arg in opts:
    if opt == '--env':
        env = arg
        break

if env == "prod":
    mysql_host = "rm-j6ctyd19275tjf7x1.mysql.rds.aliyuncs.com"
    mysql_user = "root"
    mysql_pwd = "qwemnb0011@"
    mysql_db = "aipainter"

    redis_host = "r-j6cmt5ayr2nq1m6zpp.redis.rds.aliyuncs.com"
    redis_db = 0
    redis_pwd = "r-j6cmt5ayr2nq1m6zpp:qpalmz1209@L"

    oss_endpoint = "oss-cn-hongkong-internal.aliyuncs.com"
else:
    mysql_host = "192.168.1.56"
    mysql_user = "root"
    mysql_pwd = "ubmy8920fe101f0aabffeeaecf1d9e04"
    mysql_db = "dl_oversea"

    redis_host = "192.168.1.56"
    redis_db = 0
    redis_pwd = "ubd83922c944c7e45b45ab5c86f98e7f"

    oss_endpoint = "oss-cn-hongkong.aliyuncs.com"

# mysql
conn = pymysql.connect(host=mysql_host, user=mysql_user, password=mysql_pwd, database=mysql_db)
cur = conn.cursor()

# redis
r = redis.StrictRedis(host=redis_host, port=6379, db=redis_db, password=redis_pwd)

# oss
ossAuth = oss2.Auth('LTAI5tMzkcDHweKw7762CzQn', '6cfufuTuoQpdwR7mckhxITwkLBWtcH')
ossBucket = oss2.Bucket(ossAuth, oss_endpoint, 'aipainter-overseas')


def writeLog(data, machine):
    now = datetime.datetime.now().strftime('%Y-%m-%d')
    path = "../logs/{}_{}.txt".format(machine, now)
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    file = open(path, 'a+')
    file.write(now + " " + data + '\n')
    file.close()


def getFileBase64(imgUrl):
    imgUrl += "!webp"
    response = requests.get(imgUrl)
    base64_data = base64.b64encode(response.content).decode('utf-8')
    return base64_data


def doTask(URL, machine):
    queueKey = "painter:{}:queue".format(machine)

    loopNum = 0

    writeLog("loop start!!!", machine)

    while loopNum < 500:
        loopNum += 1

        sdinput = r.lpop(queueKey)
        if sdinput == None:
            writeLog("no task!!", machine)
            cur.execute('select 1')
            time.sleep(1)
            continue
        toStr = str(sdinput, encoding="utf-8")
        writeLog(toStr, machine)
        arr = json.loads(toStr)

        task_id = arr['task_id']
        writeLog("task_id : {} start!".format(task_id), machine)

        taskKey = "painter:task:{}".format(task_id)
        curUserId = arr['user_id'] if arr['user_id'] > 0 else arr["device_id"]
        userLenKey = "painter:user:len:{}".format(curUserId)
        pngFile = "{}.png".format(task_id)
        webpFile = "{}.webp".format(task_id)
        headers = {"User-Agent": arr["model"]}
        # drawing_image = arr['drawing_url']
        isHigh = arr['isHigh']
        style_id = arr["style_id"]

        # 默认值drawing_image
        if "drawing_url" not in arr:
            drawing_image = ""
        else:
            drawing_image = arr['drawing_url']

        # 是否是延时任务
        delay_task = False
        if "is_dealy" in arr and arr["is_delay"] == 1:
            delay_task = True

        # 设置处理中状态
        r.hset(taskKey, 'status', 1)
        try:
            if arr['type'] == 1 or arr['type'] == 2:
                writeLog("任务类型：txt2img", machine)
                sampler_index = "Euler a" if arr["type"] == 1 else "DPM++ 2M Karras"
                payload = {
                    "prompt": arr["translator_prompt"],
                    "steps": 20,
                    "cfg_scale": "7.0",
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt": arr["negative_prompt"],
                    "sampler_index": sampler_index,
                }
                if isHigh == 1:
                    hig_payload = {
                        "enable_hr": True,
                        "denoising_strength": 0.5,
                        "hr_scale": 2,
                        "hr_upscaler": "Latent"
                    }
                    payload.update(hig_payload)

                payload_json = json.dumps(payload)

                response = requests.post(url=f'{URL}/sdapi/v1/txt2img', data=payload_json, headers=headers,
                                         timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)

            else:
                writeLog("任务类型：img2img", machine)

                base64Data = getFileBase64(drawing_image)

                rand_int = random.randint(0, 1)
                interrogateModel = ""
                if rand_int == 0:
                    interrogateModel = "clip"
                else:
                    interrogateModel = "deepdanbooru"

                writeLog("interrogateModel：{}".format(interrogateModel), machine)

                payload = {
                    "image": base64Data,
                    "model": interrogateModel
                    # deepdanbooru,clip
                }

                payload_json_pro = json.dumps(payload)
                # writeLog("payload_json_pro：{}，header:{}".format(payload_json_pro,headers),machine)
                responsePro = requests.post(url=f'{URL}/sdapi/v1/interrogate', data=payload_json_pro, headers=headers,
                                            timeout=50).json()

                if "caption" not in responsePro:
                    responsePro['caption'] = ""
                arr["prompt"] = responsePro['caption']
                writeLog("图片解析的描述词：{}".format(responsePro['caption']), machine)

                base64dataArray = []
                base64dataArray.append(base64Data)

                payload = {
                    "init_images": base64dataArray,
                    "prompt": responsePro['caption'],
                    "steps": 20,
                    "denoising_strength": 0.65,
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt": arr["negative_prompt"],
                    "sampler_index": "DPM++ 2M Karras",
                    "restore_faces": True
                }

                payload_json = json.dumps(payload)
                response = requests.post(url=f'{URL}/sdapi/v1/img2img', data=payload_json, headers=headers,
                                         timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)

            # png转webp Start
            # sCmd = "ffmpeg -nostdin -i " + pngFile + "  " + webpFile
            sCmd = "ffmpeg -nostdin -i " + pngFile + " -c:v libwebp -quality 90 -compression_level 6 " + webpFile
            writeLog("图片转换: {}".format(sCmd), machine)
            os.system(sCmd)
            writeLog("png转webp成功", machine)
            # png转webp End

            # 上传图片&删除本地图片
            path = "sd/opus/%s/%s" % (time.strftime("%Y-%m-%d", time.localtime()), webpFile)
            writeLog("本地图片地址：{}".format(webpFile), machine)
            writeLog("oss图片地址：{}".format(path), machine)
            ossBucket.put_object_from_file(path, webpFile)

            # 写入数据&更新状态
            url = 'https://overseas-files.dalleai.art/' + path

            sql = """
            insert into ap_opus(user_id, prompt, prompt_use, task_id, url, height, width, negative_prompt, type, drawing_id, drawing_image, is_new,device_id, is_high, dot, style_id, is_style_orginal, platform, is_try, machine_id)
            value(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            # writeLog(sql % (arr['user_id'], arr['prompt'], arr['task_id'], url, arr['height'], arr['width'], arr['prompt'], arr['negative_prompt'], arr['type'],arr['model_id'],drawing_image,1,arr["device_id"]),machine)
            cur.execute(sql, (
            arr['user_id'], arr['prompt'], arr['translator_prompt'], arr['task_id'], url, arr['height'], arr['width'],
            arr['negative_prompt'], arr['type'], arr['model_id'], drawing_image, 1, arr["device_id"], isHigh,
            arr["dot"], style_id, arr["is_style_orginal"], arr["platform"], arr["is_try"], machine))

            sql = "update ap_model_list set use_num =  use_num + 1 where id = %s"
            cur.execute(sql, (arr['model_id']))

            if style_id > 0:
                sql = "update ap_model_style set use_num =  use_num + 1 where id = %s"
                cur.execute(sql, (style_id))

            conn.commit()

            # if arr['type'] == 3 :
            #    upload_obj = urlparse(drawing_image)
            #    upload_path = upload_obj.path[1:]
            #    ossBucket.delete_object(upload_path)
            #    writeLog("删除用户上传图片：{}".format(upload_path),machine)

            if not delay_task:
                r.hset(taskKey, 'status', 2)
                r.hset(taskKey, 'url', url)
                r.hset(taskKey, 'prompt', arr['prompt'])
            r.decrby(userLenKey, 1)
            r.expire(userLenKey, 600)
            writeLog("task_id : {} complate!".format(task_id), machine)

        except Exception as e:
            if not delay_task:
                r.hset(taskKey, 'status', 0)
            # r.rpush(queueKey, sdinput)
            writeLog(traceback.format_exc(), machine)

            # fail_amount 失败次数
            if ("fail_amount" not in arr):
                arr["fail_amount"] = 1
                t_sdinput = json.dumps(arr)
                r.rpush(queueKey, t_sdinput)
            else:
                url = 'https://overseas-files.dalleai.art/default.png'
                if int(arr['dot']) > 0:
                    sql = "update ap_user set dot = dot + %s where id = %s"
                    cur.execute(sql, (arr['dot'], arr['user_id']))

                    sql = "insert ap_dot_log(user_id, value, type) value(%s, %s, %s)"
                    cur.execute(sql, (arr['user_id'], arr['dot'], 5))

                    url = 'https://overseas-files.dalleai.art/default-dot.png'

                r.decrby(userLenKey, 1)
                r.expire(userLenKey, 600)

                writeLog("task : {} fail!".format(sdinput), machine)


        finally:
            if os.path.exists(pngFile):
                os.remove(pngFile)
            if os.path.exists(webpFile):
                os.remove(webpFile)

    writeLog("loop stop! have a rest!!!", machine)
